import torch
import cv2 as cv
from model.yolov5 import YoloV5
from model.decode import decode_box, scale_coords, plot_one_box
from utils.utils import letter_box_ractangle


model = YoloV5(classify=True)
model.load_state_dict(torch.load('yolov5.pth'))
# image_origin = cv.imread('E:/dataset/hand/6/IMG_5059.jpg')
image_origin = cv.imread('hand3.png')
assert image_origin is not None, 'Image does not exit.'
image = letter_box_ractangle(image_origin, 128)
image = image.transpose(2, 0, 1)/256
print('Network loading complete.')
model.eval()
with torch.no_grad():
    image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
    predict = model(image).squeeze()
conf, label = torch.max(predict, 0)
conf = torch.sigmoid(conf)
print(f'{int(label)}:{conf:.4f}')
# cv.imwrite('test.png', image_origin)
